#include "share/HamiForm.h"

clusterHamiltonian * hm;

int 	** HAMI_FORM_LIST;
int		HAMI_FORM_LEN;
int		** CLUSTER_2_HamiFORM;

int cmp(const void * a, const void * b);
void 	hamiform_statistics();
int 	test_form_exist(int * form_stat);
int		Get_Sgl_State_Dim(int idx, int * sgl_state_dim);

void generate_hamiltonian_forms()
{
	int i;
    hamiform_statistics();

    hm = new clusterHamiltonian[HAMI_FORM_LEN];

    int sgl_state_dim [MAX_CLST_ORDER];

    int hamiform_n_spin;
    for(i = 0; i < HAMI_FORM_LEN; i ++)
    {
    	hamiform_n_spin = Get_Sgl_State_Dim(i, sgl_state_dim);
    	hm[i].CreateHamiForm(hamiform_n_spin, sgl_state_dim, IS_BLK);
    }
}

int clusterHamiltonian::CreateHamiForm(int nspin, int * sgl_dim, int isBLK)
{
    int i;
    initialize(nspin, sgl_dim, isBLK);
    cout << "Creating Hamiltonian Spin= " << spin_num << "; full dim = " << full_dim << "; blk no. = " << nBlock << endl;
    for(i = 0; i < nBlock; i++)
    {
        cout << "\t BLK " << i << ": dim = " << dim[i] << ", " << dim_start[i] << " ~ " << dim_end[i] << endl;
        createHami(dim_start[i], dim_end[i], matList[i]);
    }

    return 0;
}

int clusterHamiltonian::initialize(int nspin, int * sgl_dim, int isBLK)
{
    int i, j, sum_excite=1;

    for(i = 0; i < nspin; i++)
        sum_excite += (sgl_dim[i] - 1);

    spin_num    = nspin;
    isBlock     = isBLK;
    nBlock      = isBlock ? sum_excite : 1 ;

    accDim = new int [spin_num];
    S_Val = new double [spin_num];

    full_dim    = 1;
    for( i = 0; i < nspin; i++ )
    {
        accDim[i] = full_dim;
        full_dim *= sgl_dim[i];
        S_Val[i] = 0.5 * ( (double) sgl_dim[i] - 1.0 );
    }

    basisList   = new basis [full_dim];
    dim_start   = new int [nBlock];
    dim_end     = new int [nBlock];
    dim         = new int [nBlock];
    if(!isBlock)
    {
        dim_start[0] = 0;
        dim_end[0] = full_dim;
        dim[0] = full_dim;
    }

    matList = new HamiElement * [nBlock];
    if(!isBlock) matList[0] =  new HamiElement [full_dim * full_dim];

//////////////////////////////////////////////////
// initialize the basis list
    int * exciteN_list = new int [spin_num];

    for(i = 0; i < full_dim; i++)
    {
        idx2excitenum(i,exciteN_list);
        basisList[i].nspin = spin_num;
        basisList[i].state = new int [spin_num];
        basisList[i].nExcite = 0;
        for(j = 0; j < spin_num; j++)
        {
            basisList[i].state[j] = exciteN_list[j];
            basisList[i].nExcite += exciteN_list[j];
        }
    }

    for(i = 0; i < full_dim; i++)
    {
    	for(j = 0; j < spin_num; j++)
    		cout << basisList[i].state[j] << "\t";
    	cout << basisList[i].nExcite << endl;
    }

    qsort(basisList, full_dim, sizeof(basis), cmp);

    cout << "after sort" << endl;
    for(i = 0; i < full_dim; i++)
    {
    	for(j = 0; j < spin_num; j++)
    		cout << basisList[i].state[j] << "\t";
    	cout << basisList[i].nExcite << endl;
    }

    int excit, ncount = 0, dim2;
    if(isBLK)
    {
    	dim_start[0]=0; excit = basisList[0].nExcite;
    	for(i = 0; i < full_dim; i++)
    	{
    		if(basisList[i].nExcite != excit)
    		{
    			dim_end[ncount] = i;
    			dim_start[ncount+1] = i;
    			dim[ncount] = dim_end[ncount] - dim_start[ncount];
    			dim2 = dim[ncount] * dim[ncount];
    			matList[ncount] = new HamiElement [ dim2 ];
//              cout <<"matList[i = " << ncount << "] is allocated!" << endl;

    			excit = basisList[i].nExcite;
    			ncount ++;
    		}
    	}
    	dim_end[ncount] = full_dim;
    	dim[ncount] = dim_end[ncount] - dim_start[ncount];
    	dim2 = dim[ncount] * dim[ncount];
    	matList[ncount] = new HamiElement [ dim2 ];
//      cout <<"matList[i = " << ncount << "] is allocated!" << endl;

    	assert(ncount + 1 == nBlock);
    }


    for(i = 0; i < nBlock; i++)
    {
        for(j = 0; j < dim[i]*dim[i]; j++)
        {
            matList[i][j].nterms_R = 0;
            matList[i][j].nterms_I = 0;
            matList[i][j].coefflist_R = NULL;
            matList[i][j].coefflist_I = NULL;
            matList[i][j].idxlist_R = NULL;
            matList[i][j].idxlist_I = NULL;
            matList[i][j].termName_R = NULL;
            matList[i][j].termName_I = NULL;
//            cout << "matList[i="<< i << "]" << "[j=" << j << "]" << "is initialised" << endl;
        }
    }

    delete [] accDim;   accDim = NULL;
    delete [] exciteN_list;
    return 0;
}

int cmp(const void * a, const void * b)
{
    return (* (basis *) b).nExcite - (* (basis *) a).nExcite;
}

void clusterHamiltonian::idx2excitenum(int idx, int* exciteN_list)
{
    int i, num;

    num = idx;
    for(i = spin_num - 1; i >= 0; i--)
    {
        exciteN_list[i] =  num / accDim[i];
        num = num - exciteN_list[i] * accDim[i];
    }
}

int clusterHamiltonian::createHami(int n_start, int n_end, HamiElement * mat)
{
    int i, j;
    int dim;

    dim = n_end - n_start;

    basis state_i, state_j;

    for(i = 0; i < dim; i ++)
    {
        state_i = basisList[n_start + i];
        for(j = 0; j <= i; j ++)
        {
            state_j = basisList[n_start + j];
            mat_elements(mat[j * dim + i], state_i, state_j);
        }

    }
    return 0;
}

int clusterHamiltonian::mat_elements(HamiElement & ele, basis state_i, basis state_j)
{
    int idx, ncount, state_diff_num;//, excit_diff
//    excit_diff = abs(state_i.nExcite - state_j.nExcite);
    
    state_diff_num = 0;
    for(idx = 0; idx < spin_num; idx ++)
    {
        if(state_i.state[idx] != state_j.state[idx])
        {
            state_diff_num ++;
        }
    }


    if(state_diff_num > 2)
    {
        ele.nterms_R = 0;
        ele.nterms_I = 0;
        return 0;
    }

/////////////////////////////////////
//else: state_diff_num <=2

    int * diff_site = new int [state_diff_num];
    ncount = 0;
    for(idx = 0; idx < spin_num; idx ++)
    {
        if(state_i.state[idx] != state_j.state[idx])
        {
            diff_site[ncount] = idx;
            ncount ++;
        }
    }



    if(state_diff_num == 0)
    {
        mat_element_diagonal(ele, state_i);
    }
    else if(state_diff_num == 1)
    {
        mat_element_off_diagonal1(ele, state_i, state_j, diff_site);
    }
    else// if(state_diff_num == 2)
    {
        mat_element_off_diagonal2(ele, state_i, state_j,diff_site);
    }


    delete [] diff_site;
    return 0;
}

int clusterHamiltonian::mat_element_diagonal(HamiElement & ele, basis state_i)
{
    int idx, state1, state2, int_term_num;

    // # of terms = terms contributed by two-body diagonal interaction: \sigma_z^i \sigma_z^j
    // and terms contributed by single body Zeeman term: \sigma_z
    int_term_num =  spin_num * (spin_num - 1) / 2;
    ele.nterms_R = spin_num + int_term_num;
    ele.nterms_I = ele.nterms_R;

    ele.coefflist_R = new double [ele.nterms_R];
    ele.coefflist_I = new double [ele.nterms_I];
    ele.termName_R = new int [ele.nterms_R];
    ele.termName_I = new int [ele.nterms_I];
    ele.idxlist_R = new int [ele.nterms_R];
    ele.idxlist_I = new int [ele.nterms_I];
    
    double sz1, sz2;
    for(idx = 0; idx < spin_num; idx++)
    {
        state1 = state_i.state[idx];
        sz1 = (double)state1 - S_Val[idx];

        ele.coefflist_R[idx] = sz1; //( state1 == 0 ? 1.0 : -1.0 );    // 0 represents spinUp; 1 represents spinDown
        ele.coefflist_I[idx] = 0.0;
        ele.termName_R[idx] = BZ_ARRAY;
        ele.termName_I[idx] = ZERO_ARRAY;
        ele.idxlist_R[idx] = idx;
        ele.idxlist_I[idx] = 0;
    }

    if(spin_num == 1) return 0;

    gsl_combination * c;

    int ncount;

    ncount = spin_num;
    int state_idx1, state_idx2;
    c = gsl_combination_calloc (spin_num, 2);
    do
    {
        state_idx1 = gsl_combination_get(c, 0);
        state_idx2 = gsl_combination_get(c, 1);
        state1 = state_i.state[state_idx1];		sz1 =  (double)state1 - S_Val[state_idx1];
        state2 = state_i.state[state_idx2];		sz2 =  (double)state2 - S_Val[state_idx2];
        ele.coefflist_R[ncount] = sz1 * sz2; //( state1 == state2 ? 1.0 : -1.0 );
        ele.coefflist_I[ncount] = 0.0;
        ele.termName_R[ncount] = ZZ_ARRAY;
        ele.termName_I[ncount] = ZERO_ARRAY;
        ele.idxlist_R[ncount] = state_idx1 * spin_num + state_idx2;
        ele.idxlist_I[ncount] = 0;
        ncount ++;
    }
    while (gsl_combination_next (c) == GSL_SUCCESS);
    gsl_combination_free (c);


    return 0;
}

int clusterHamiltonian::mat_element_off_diagonal1(HamiElement & ele, basis state_i, basis state_j, int * diff_site)
{
    int i, diff_idx, diff_state_i, diff_state_j, comm_state, delta_state;
    double sz_common, diff_sign, coeff, Sval, mi_val, mj_val, Sval_common;

    diff_idx = diff_site[0];
    Sval = S_Val[diff_idx];
    diff_state_i = state_i.state[diff_idx];
    diff_state_j = state_j.state[diff_idx];
    delta_state = diff_state_i - diff_state_j;

//    diff_state = state_j.state[diff_idx];
//    diff_sign = ( diff_state == 0 ? 1.0 : -1.0 );

    if( delta_state == 1 || delta_state == -1)
    {
    	mi_val = ((double) diff_state_i) -Sval;
    	mj_val = ((double) diff_state_j) -Sval;
    	coeff = 0.5 * sqrt(Sval * (Sval + 1) - mi_val * mj_val);
    	diff_sign = (delta_state == 1 ? -1.0 : 1.0);

    	ele.nterms_R = spin_num;
    	ele.nterms_I = spin_num;

    	ele.coefflist_R = new double [spin_num ];
    	ele.coefflist_I = new double [spin_num ];
    	ele.termName_R = new int [spin_num ];
    	ele.termName_I = new int [spin_num ];
    	ele.idxlist_R = new int [spin_num ];
    	ele.idxlist_I = new int [spin_num ];


    	for(i = 0; i < spin_num; i++)
    	{
    		if(i < diff_idx)
    		{
    			Sval_common = S_Val[i];
    			comm_state = state_j.state[i];
        		sz_common = ((double)comm_state) - Sval_common;

        		ele.coefflist_R[i] = sz_common * coeff;
    			ele.coefflist_I[i] = sz_common * diff_sign * coeff;
    			ele.termName_R[i] = ZX_ARRAY;
    			ele.termName_I[i] = ZY_ARRAY;
    			ele.idxlist_R[i] = i * spin_num + diff_idx;
    			ele.idxlist_I[i] = i * spin_num + diff_idx;
    		}
    		if(i == diff_idx)
    		{
    			ele.coefflist_R[i] = coeff;
    			ele.coefflist_I[i] = diff_sign * coeff;
    			ele.termName_R[i] = BX_ARRAY;
    			ele.termName_I[i] = BY_ARRAY;
    			ele.idxlist_R[i] = i;
    			ele.idxlist_I[i] = i;
    		}
    		if(i > diff_idx)
    		{
    			Sval_common = S_Val[i];
    			comm_state = state_j.state[i];
        		sz_common = ((double)comm_state) - Sval_common;

        		ele.coefflist_R[i] = sz_common * coeff;
    			ele.coefflist_I[i] = sz_common * diff_sign * coeff;
    			ele.termName_R[i] = XZ_ARRAY;
    			ele.termName_I[i] = YZ_ARRAY;
    			ele.idxlist_R[i] = i * spin_num + diff_idx;
    			ele.idxlist_I[i] = i * spin_num + diff_idx;
    		}
    	}
    }
    else
    {
    	ele.nterms_R = 0;
    	ele.nterms_I = 0;
    }

    return 0;
}

int clusterHamiltonian::mat_element_off_diagonal2(HamiElement & ele, basis state_i, basis state_j, int * diff_site)
{
    int diff_idx1, diff_idx2, polar;
    int state_i_1, state_i_2, state_j_1, state_j_2, delta_1, delta_2;
    double Sval1, Sval2, mi_1, mi_2, mj_1, mj_2, coeff1, coeff2;

    diff_idx1 = diff_site[0];
    diff_idx2 = diff_site[1];
    assert(diff_idx1 < diff_idx2);

    state_i_1 = state_i.state[diff_idx1];
    state_i_2 = state_i.state[diff_idx2];
    state_j_1 = state_j.state[diff_idx1];
    state_j_2 = state_j.state[diff_idx2];
    delta_1 = state_i_1 - state_j_1;
    delta_2 = state_i_2 - state_j_2;

    Sval1 = S_Val[diff_idx1];	mi_1 = (double)state_i_1 - Sval1;	mj_1 = (double)state_j_1 - Sval1;
    Sval2 = S_Val[diff_idx2];	mi_2 = (double)state_i_2 - Sval2;	mj_2 = (double)state_j_2 - Sval2;
	coeff1 = 0.5 * sqrt(Sval1 * (Sval1 + 1) - mi_1 * mj_1);
	coeff2 = 0.5 * sqrt(Sval2 * (Sval2 + 1) - mi_2 * mj_2);

	ele.coefflist_R = new double [ 2 ];
	ele.coefflist_I = new double [ 2 ];
	ele.termName_R = new int [ 2 ];
	ele.termName_I = new int [ 2 ];
	ele.idxlist_R = new int [ 2 ];
	ele.idxlist_I = new int [ 2 ];

	if(delta_1 == 1 && delta_2 == 1)
    {
    	ele.nterms_R = 2;
    	ele.nterms_I = 2;

    	ele.coefflist_R[0] = coeff1 * coeff2;                   ele.coefflist_I[0] = - coeff1 * coeff2;
    	ele.termName_R[0]  = XX_ARRAY;                          ele.termName_I[0] = XY_ARRAY;
    	ele.idxlist_R[0]   = diff_idx1 * spin_num + diff_idx2;  ele.idxlist_I[0] = diff_idx1 * spin_num + diff_idx2;

    	ele.coefflist_R[1] = - coeff1 * coeff2;					ele.coefflist_I[1] = - coeff1 * coeff2;
    	ele.termName_R[1]  = YY_ARRAY;							ele.termName_I[1] = YX_ARRAY;
    	ele.idxlist_R[1]   = diff_idx1 * spin_num + diff_idx2;	ele.idxlist_I[1] = diff_idx1 * spin_num + diff_idx2;
    }
    else if(delta_1 == 1 && delta_2 == -1)
    {
    	ele.nterms_R = 2;
    	ele.nterms_I = 2;

    	ele.coefflist_R[0] = coeff1 * coeff2;					ele.coefflist_I[0] = coeff1 * coeff2;
    	ele.coefflist_R[1] = coeff1 * coeff2;					ele.coefflist_I[1] = - coeff1 * coeff2;


    	ele.termName_R[0] = XX_ARRAY;							ele.termName_I[0] = XY_ARRAY;
    	ele.termName_R[1] = YY_ARRAY;							ele.termName_I[1] = YX_ARRAY;

    	ele.idxlist_R[0] = diff_idx1 * spin_num + diff_idx2;	ele.idxlist_I[0] = diff_idx1 * spin_num + diff_idx2;
    	ele.idxlist_R[1] = diff_idx1 * spin_num + diff_idx2;	ele.idxlist_I[1] = diff_idx1 * spin_num + diff_idx2;
    }
    else if(delta_1 == -1 && delta_2 == 1)
    {
    	ele.nterms_R = 2;
    	ele.nterms_I = 2;

    	ele.coefflist_R[0] = coeff1 * coeff2;                         ele.coefflist_I[0] = - coeff1 * coeff2;
    	ele.coefflist_R[1] = coeff1 * coeff2;                         ele.coefflist_I[1] = coeff1 * coeff2;


    	ele.termName_R[0] = XX_ARRAY;                                 ele.termName_I[0] = XY_ARRAY;
    	ele.termName_R[1] = YY_ARRAY;                                 ele.termName_I[1] = YX_ARRAY;

    	ele.idxlist_R[0] = diff_idx1 * spin_num + diff_idx2;          ele.idxlist_I[0] = diff_idx1 * spin_num + diff_idx2;
    	ele.idxlist_R[1] = diff_idx1 * spin_num + diff_idx2;          ele.idxlist_I[1] = diff_idx1 * spin_num + diff_idx2;
    }
    else if(delta_1 == -1 && delta_2 == -1)
    {
    	ele.nterms_R = 2;
    	ele.nterms_I = 2;

    	ele.coefflist_R[0] = coeff1 * coeff2;                   ele.coefflist_I[0] = coeff1 * coeff2;
    	ele.termName_R[0]  = XX_ARRAY;                          ele.termName_I[0] = XY_ARRAY;
    	ele.idxlist_R[0]   = diff_idx1 * spin_num + diff_idx2;  ele.idxlist_I[0] = diff_idx1 * spin_num + diff_idx2;

    	ele.coefflist_R[1] = - coeff1 * coeff2;                 ele.coefflist_I[1] = coeff1 * coeff2;
    	ele.termName_R[1]  = YY_ARRAY;                          ele.termName_I[1] = YX_ARRAY;
    	ele.idxlist_R[1]   = diff_idx1 * spin_num + diff_idx2;  ele.idxlist_I[1] = diff_idx1 * spin_num + diff_idx2;
    }
    else
    {
    	ele.nterms_R = 0;
    	ele.nterms_I = 0;
    }

    return 0;
}

double clusterHamiltonian::ElementVal_R(double** interaction, int blk, int idx)
{
    int i, nterms, interaction_name, interaction_seq;
    double res = 0.0,  coeff;

    nterms = matList[blk][idx].nterms_R;

    for(i = 0; i < nterms; i++)
    {
        coeff               = matList[blk][idx].coefflist_R[i];
        interaction_name    = matList[blk][idx].termName_R[i];
        interaction_seq     = matList[blk][idx].idxlist_R[i];
        res +=  coeff * interaction[interaction_name][interaction_seq];
    }

    return(res);
}

double clusterHamiltonian::ElementVal_I(double** interaction, int blk, int idx)
{
    int i, nterms, interaction_name, interaction_seq;
    double res = 0.0,  coeff;

    nterms = matList[blk][idx].nterms_I;

    for(i = 0; i < nterms; i++)
    {
        coeff               = matList[blk][idx].coefflist_I[i];
        interaction_name    = matList[blk][idx].termName_I[i];
        interaction_seq     = matList[blk][idx].idxlist_I[i];
        res +=  coeff * interaction[interaction_name][interaction_seq];
    }

    return(res);
}

void hamiform_statistics()
{
	int i, ii, j, q, clst_count;

	int spin_val, idx;

	HAMI_FORM_LIST = new int * [MAX_HAMI_FORM];
	HAMI_FORM_LEN = 0;

	CLUSTER_2_HamiFORM = new int * [CALCULATED_ORDER];

	int * form_stat = new int [MAX_SPECIES + 2];	// form_stat[0] is the species number; form_stat[1] is the number of the spins
	for(i = 0; i < MAX_SPECIES + 2; i++)
		form_stat[i] = 0;

	clst_count = 0;
	for(i = 0; i < CALCULATED_ORDER; i ++)
	{
		CLUSTER_2_HamiFORM[i] = new int [NUM_CLST[i]];

		for(j = 0; j < NUM_CLST[i]; j++)
		{
			for(ii = 0; ii < MAX_SPECIES + 2; ii++)
				form_stat[ii] = 0;

			for (q = 0; q < i + 1; q++ )
			{
				idx = CLST_INDEX[i][j*(i+1) + q];
				spin_val = ATOM_SPIN[idx];
				if(form_stat[ spin_val ] == 0)
					form_stat[0] ++;
				form_stat[1]++;
				form_stat[ spin_val ] ++;
			}
//			cout << i << ": ";
//			for(ii = 0; ii < MAX_SPECIES + 2; ii++)
//				 cout << form_stat[ii] << ", ";
//			cout << endl;

			CLUSTER_2_HamiFORM[i][j] = test_form_exist(form_stat);
			clst_count ++;
		}
	}
	assert(clst_count == TOTALCLST);
	delete [] form_stat;
}

int test_form_exist(int * form_stat)
{
	int res;
	int i, j;
	int isExist = 0;	// not exist by default
	int isSame;

	for(i = 0; i < HAMI_FORM_LEN; i++)
	{
		isSame = 1;
		for(j = 0; j < MAX_SPECIES + 2; j ++)
		{
			if(form_stat[j] != HAMI_FORM_LIST[i][j])
			{
				isSame = 0;
				break;
			}
		}
		if(isSame == 1)
		{
			isExist = 1;
			res = i;
			break;
		}
	}
	if(isExist == 0)
	{
		HAMI_FORM_LIST[HAMI_FORM_LEN] = new int [MAX_SPECIES + 2];
		for(j = 0; j < MAX_SPECIES + 2; j ++)
			HAMI_FORM_LIST[HAMI_FORM_LEN][j] = form_stat[j];
		res = HAMI_FORM_LEN;
		HAMI_FORM_LEN ++;
	}
	return res;
}

int Get_Sgl_State_Dim(int idx, int * sgl_state_dim)
{
	int i, j, count, nspin;
	for(i = 0; i < MAX_CLST_ORDER; i ++)
		sgl_state_dim[i] = 0;

	nspin = HAMI_FORM_LIST[idx][1];
	for(i = 0; i < MAX_SPECIES + 2; i++)
		cout << HAMI_FORM_LIST[idx][i] << ", " ;
	cout << endl;

	count = 0;
	for(i = 2; i < MAX_SPECIES + 2; i++)
	{
		for(j = 0; j < HAMI_FORM_LIST[idx][i]; j++)
		{
			sgl_state_dim[count] = i;
			count ++;
		}
	}
	cout << ">>>>";
	for(i = 0; i < MAX_CLST_ORDER; i ++)
		cout << sgl_state_dim[i] << ", ";
	cout << endl;

	return nspin;
}
